import os
import logging
import re
import time
import asyncio
import concurrent.futures
# 移除了 hashlib 和 threading 导入，因为不再使用缓存机制
from typing import Dict, Any, List, Tuple, Optional
from pymilvus import MilvusClient
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from sqlalchemy import create_engine, text
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage

# 导入配置 - 支持相对导入和绝对导入
try:
    # 尝试相对导入（作为包模块运行时）
    from .config import (
        MYSQL_DB_URL, MILVUS_HOST, MILVUS_PORT,
        OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, OPENAI_EMBEDDING_MODEL,
        AGENT_MAX_RETRIES, AGENT_TIMEOUT, LOG_LEVEL
    )
except ImportError:
    # 回退到绝对导入（直接运行脚本时）
    import sys
    import os
    # 添加当前目录到Python路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    if current_dir not in sys.path:
        sys.path.insert(0, current_dir)

    from config import (
        MYSQL_DB_URL, MILVUS_HOST, MILVUS_PORT,
        OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, OPENAI_EMBEDDING_MODEL,
        AGENT_MAX_RETRIES, AGENT_TIMEOUT, LOG_LEVEL
    )

# 环境配置
logging.basicConfig(level=getattr(logging, LOG_LEVEL.upper()),
                   format='%(asctime)s - %(levelname)s - %(message)s')

# 优化配置已移至 optimization_solution.py

class Text2SQLAgent:
    def __init__(self):
        # 初始化组件
        self.model = ChatOpenAI(
            model=OPENAI_MODEL,
            api_key=OPENAI_API_KEY,
            base_url=OPENAI_BASE_URL,
        )

        # 初始化嵌入函数
        try:
            self.embedding_fn = OpenAIEmbeddings(
                api_key=OPENAI_API_KEY,
                base_url=OPENAI_BASE_URL,
                model=OPENAI_EMBEDDING_MODEL
            )
            logging.info("嵌入函数初始化成功")

        except Exception as e:
            logging.error(f"嵌入函数初始化失败：{e}")
            raise

        # 🚀 统一嵌入生成优化：一次生成，多次复用
        # 移除了复杂的缓存机制，专注于核心的统一嵌入优化

        # 初始化Milvus客户端
        try:
            self.milvus_client = MilvusClient(
                host=MILVUS_HOST,
                port=MILVUS_PORT
            )
            logging.info("Milvus客户端初始化成功")
        except Exception as e:
            logging.error(f"Milvus客户端初始化失败：{e}")
            raise

        # 初始化数据库连接
        try:
            self.engine = create_engine(MYSQL_DB_URL)
            logging.info("数据库连接初始化成功")
        except Exception as e:
            logging.error(f"数据库连接初始化失败：{e}")
            raise
        
        # 创建工具列表
        self.tools = self._create_tools()

        # 创建智能体
        self.agent = create_react_agent(
            model=self.model,
            tools=self.tools,
            name="text2sql_expert",
            debug=True,
            prompt="""你是一个专业的电影数据库查询助手，专门帮助用户查找和分析电影相关信息。

你的专业领域包括：
• 电影信息查询：电影详情、演员资料、导演信息
• 票房统计分析：收入数据、排行榜、趋势分析
• 电影分类查询：按类型、年份、评分等筛选
• 租赁数据分析：客户信息、观影记录、租赁统计

请专注于电影数据库相关的查询和分析，为用户提供准确、详细的电影信息服务。

对于其他查询，你的工作流程：
1. 如果遇到嵌入相关错误，首先使用diagnose_embedding_tool诊断问题
2. 优先使用smart_query_tool进行一站式查询
3. 如果smart_query_tool不可用，则按以下步骤：
   - 使用concurrent_retrieve_all_tool获取相关信息
   - 使用generate_sql_tool生成查询语句
   - 使用execute_sql_tool执行查询
   - 使用convert_result_to_natural_language_tool转换结果
4. 如果查询失败，使用fix_sql_tool修复后重试

关键规则：
- 🚫 绝不能提及SQL、数据库、工具调用等技术术语
- ✅ 所有回复都用自然、友好的语言
- ✅ 专注于用户想要的信息和数据
- ✅ 如果查询失败，提供友好的建议

请严格按照这个流程工作，确保：
- 优先使用smart_query_tool提升性能
- 最终提供易懂的自然语言结果
- 遇到错误时能够自动修复并重试
- 所有回复都对用户友好，避免技术术语
- 突出用户关心的关键信息和数据"""
        )
        
        # 创建状态图
        self.workflow = self.create_workflow()
    
    def get_agent(self):
        """获取智能体实例"""
        return self.agent
    
    def get_tools(self):
        """获取工具列表"""
        return self.tools
    
    def get_model(self):
        """获取模型实例"""
        return self.model

    def _preload_collections(self):
        """预加载Milvus集合以提高性能"""
        collections = ["ddl_knowledge", "q2sql_knowledge", "dbdesc_knowledge"]
        for collection in collections:
            try:
                self.milvus_client.load_collection(collection)
                logging.info(f"成功预加载集合：{collection}")
            except Exception as e:
                logging.warning(f"预加载集合 {collection} 失败：{e}")
                # 不抛出异常，允许运行时加载

    def generate_unified_embedding(self, question: str) -> List[float]:
        """
        🚀 统一嵌入生成 - 一次生成，多次复用

        简化版本：直接生成嵌入向量，不使用缓存机制
        核心优化：同一个问题只生成一次嵌入向量，然后在所有检索任务中复用

        Args:
            question: 用户问题

        Returns:
            嵌入向量
        """
        # 标准化问题文本
        normalized_question = question.strip().lower()

        # 生成嵌入向量
        logging.info(f"🔄 生成统一嵌入: {question[:30]}...")
        start_time = time.time()

        # 🎯 关键：只调用1次嵌入API
        unified_embedding = self.embedding_fn.embed_query(normalized_question)

        generation_time = time.time() - start_time
        logging.info(f"✅ 统一嵌入生成完成: {generation_time:.3f}s")

        return unified_embedding

    def _concurrent_retrieve_single_with_unified_embedding(self, collection_name: str,
                                                        unified_embedding: List[float],
                                                        limit: int, output_fields: List[str]) -> Tuple[str, Any]:
        """
        🚀 使用统一嵌入的单个检索操作（优化版）

        这个方法不再重复生成嵌入，而是复用传入的统一嵌入向量

        Args:
            collection_name: 集合名称
            unified_embedding: 统一生成的嵌入向量（复用）
            limit: 检索数量限制
            output_fields: 输出字段

        Returns:
            (集合名称, 检索结果)
        """
        try:
            # 🎯 关键优化：直接使用统一嵌入，不再重复生成
            logging.debug(f"🔍 使用统一嵌入检索: {collection_name}")

            # 执行向量搜索
            results = self.milvus_client.search(
                collection_name=collection_name,
                data=[unified_embedding],  # 复用统一嵌入向量
                limit=limit,
                output_fields=output_fields
            )

            if results and len(results) > 0:
                logging.debug(f"✅ 检索成功: {collection_name}, 结果数: {len(results[0])}")
                return collection_name, results[0]
            else:
                logging.debug(f"⚠️ 检索无结果: {collection_name}")
                return collection_name, []

        except Exception as e:
            logging.error(f"❌ 统一嵌入检索{collection_name}失败: {str(e)}")
            return collection_name, None

    def _concurrent_retrieve_single(self, collection_name: str, question: str,
                                  limit: int, output_fields: List[str]) -> Tuple[str, Any]:
        """
        单个检索操作（保持向后兼容）

        注意：这个方法仍然保留，用于向后兼容，但建议使用统一嵌入版本
        """
        try:
            # 生成嵌入向量（旧方式，每次都调用API）
            q_emb = self.embedding_fn.embed_query(question.strip())

            # 执行搜索
            results = self.milvus_client.search(
                collection_name=collection_name,
                data=[q_emb],
                limit=limit,
                output_fields=output_fields
            )

            if results and len(results) > 0:
                return collection_name, results[0]
            else:
                return collection_name, []

        except Exception as e:
            logging.error(f"并发检索{collection_name}时出错：{str(e)}")
            return collection_name, None

    def concurrent_retrieve_all(self, question: str, max_workers: int = 3) -> Dict[str, Any]:
        """
        🚀 优化版并发检索 - 集成统一嵌入生成

        核心优化：
        1. 一次嵌入生成，多次向量复用
        2. 真正的并发检索（不再重复调用嵌入API）
        3. 智能缓存机制
        4. 保持向后兼容的接口

        Args:
            question: 用户问题
            max_workers: 最大并发工作线程数

        Returns:
            包含所有检索结果的字典，新增性能统计信息
        """
        start_time = time.time()

        if not question or not question.strip():
            return {"error": "问题不能为空"}

        try:
            logging.info(f"🚀 开始统一嵌入并发检索: {question[:50]}...")

            # 🎯 步骤1：统一嵌入生成（只调用1次API）
            unified_embedding = self.generate_unified_embedding(question)

            if unified_embedding is None:
                logging.error("❌ 统一嵌入生成失败")
                return {"error": "嵌入生成失败"}

            logging.info(f"✅ 统一嵌入生成完成")

            # 🚀 步骤2：定义检索任务配置（使用统一嵌入）
            retrieval_tasks = [
                ("ddl_knowledge", 3, ["ddl_text"]),
                ("q2sql_knowledge", 3, ["question", "sql_text"]),
                ("dbdesc_knowledge", 5, ["table_name", "column_name", "description"])
            ]

            results = {}

            # 🚀 步骤3：使用统一嵌入进行真正的并发检索
            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                # 提交所有检索任务（复用统一嵌入）
                future_to_collection = {
                    executor.submit(
                        self._concurrent_retrieve_single_with_unified_embedding,
                        collection_name,
                        unified_embedding,  # 🎯 复用统一嵌入向量
                        limit,
                        output_fields
                    ): collection_name
                    for collection_name, limit, output_fields in retrieval_tasks
                }

                # 收集结果
                for future in concurrent.futures.as_completed(future_to_collection):
                    collection_name = future_to_collection[future]
                    try:
                        result_collection, result_data = future.result()
                        results[result_collection] = result_data
                        logging.debug(f"✅ 检索完成: {result_collection}")
                    except Exception as e:
                        logging.error(f"❌ 检索{collection_name}失败: {e}")
                        results[collection_name] = None

            total_time = time.time() - start_time

            # 🎯 步骤4：添加性能统计信息
            results["_performance_stats"] = {
                "total_time": total_time,
                "api_calls_saved": 2,  # 相比原来3次调用，节省2次
                "optimization_applied": True,
                "unified_embedding_used": True
            }

            logging.info(f"🎉 统一嵌入并发检索完成: {total_time:.3f}s, API节省: 2次")

            return results

        except Exception as e:
            logging.error(f"❌ 统一嵌入并发检索失败: {str(e)}")
            return {
                "error": str(e),
                "_performance_stats": {
                    "total_time": time.time() - start_time,
                    "optimization_applied": False,
                    "unified_embedding_used": False
                }
            }

    # 移除了 get_embedding_performance_stats 和 clear_embedding_cache 方法
    # 简化版本不再需要性能统计和缓存管理

    def format_concurrent_results(self, results: Dict[str, Any]) -> Dict[str, str]:
        """格式化并发检索结果"""
        formatted = {}

        # 格式化DDL结果
        if "ddl_knowledge" in results and results["ddl_knowledge"]:
            ddl_hits = results["ddl_knowledge"]
            ddl_context = "\n".join(hit.get("ddl_text", "") for hit in ddl_hits)
            formatted["ddl"] = f"检索到的数据库结构信息：\n{ddl_context}"
        else:
            formatted["ddl"] = "没有找到相关的数据库结构信息"

        # 格式化示例结果
        if "q2sql_knowledge" in results and results["q2sql_knowledge"]:
            example_hits = results["q2sql_knowledge"]
            examples = "\n".join(
                f"问题：{hit.get('question', '')}\nSQL：{hit.get('sql_text', '')}\n"
                for hit in example_hits
            )
            formatted["examples"] = f"检索到的相似示例：\n{examples}"
        else:
            formatted["examples"] = "没有找到相关的查询示例"

        # 格式化描述结果
        if "dbdesc_knowledge" in results and results["dbdesc_knowledge"]:
            desc_hits = results["dbdesc_knowledge"]
            descriptions = "\n".join(
                f"{hit.get('table_name', '')}.{hit.get('column_name', '')}: {hit.get('description', '')}"
                for hit in desc_hits
            )
            formatted["descriptions"] = f"检索到的字段描述：\n{descriptions}"
        else:
            formatted["descriptions"] = "没有找到相关的字段描述"

        return formatted
    
    def _create_tools(self):
        """创建工具列表"""
        from langchain_core.tools import tool

        # 预加载Milvus集合以提高性能
        self._preload_collections()



        @tool
        def concurrent_retrieve_all_tool(question: str) -> str:
            """
            🚀 优化版并发检索工具 - 集成统一嵌入生成

            并发检索所有相关信息（DDL、示例、字段描述）
            使用统一嵌入生成优化，减少API调用次数
            """
            try:
                # 执行优化后的并发检索
                start_time = time.time()
                results = self.concurrent_retrieve_all(question)
                end_time = time.time()

                if "error" in results:
                    return f"❌ 检索失败: {results['error']}"

                # 格式化结果
                formatted = self.format_concurrent_results(results)

                # 获取性能统计信息
                perf_stats = results.get("_performance_stats", {})

                # 构建性能信息字符串
                perf_info = ""
                if perf_stats.get("optimization_applied"):
                    api_saved = perf_stats.get("api_calls_saved", 0)
                    perf_info = f"""
🚀 性能优化信息：
   - 统一嵌入: 已生成
   - API调用节省: {api_saved}次
   - 优化生效: ✅"""

                # 合并所有检索结果
                combined_result = f"""✅ 并发检索完成（耗时：{end_time - start_time:.2f}秒）{perf_info}

{formatted['ddl']}

{formatted['examples']}

{formatted['descriptions']}"""

                return combined_result

            except Exception as e:
                logging.error(f"❌ 并发检索失败：{str(e)}")
                return f"❌ 并发检索失败：{str(e)}"

        # 移除了 get_embedding_performance_tool 工具方法
        # 简化版本不再提供性能统计功能

        @tool
        def generate_sql_tool(context: str, question: str) -> str:
            """基于上下文和问题生成SQL语句"""
            prompt = f"""基于以下信息生成SQL语句：

{context}

用户问题：{question}

请只返回SQL语句，不要包含任何解释或说明。"""

            response = self.model.invoke(prompt)
            sql = self.extract_sql(response.content)
            return f"生成的SQL语句：{sql}"

        @tool
        def execute_sql_tool(sql: str) -> str:
            """执行SQL语句并返回结构化结果"""
            try:
                # 清理SQL语句
                clean_sql = sql.strip().rstrip(';')
                if not clean_sql:
                    return "SQL语句为空"

                with self.engine.connect() as conn:
                    result = conn.execute(text(clean_sql))
                    cols = list(result.keys())
                    rows = result.fetchall()

                    # 转换为字典列表格式，便于后续处理
                    data = []
                    for row in rows:
                        row_dict = {}
                        for i, col in enumerate(cols):
                            # 处理不同数据类型
                            value = row[i]
                            if value is None:
                                row_dict[col] = None
                            elif isinstance(value, (int, float, str, bool)):
                                row_dict[col] = value
                            else:
                                row_dict[col] = str(value)
                        data.append(row_dict)

                    # 返回结构化结果
                    return f"SQL执行成功！\n查询结果：{len(data)}行数据\n列名：{cols}\n数据：{data}"

            except Exception as e:
                error_msg = str(e)
                logging.error(f"SQL执行失败：{error_msg}")
                return f"SQL执行失败：{error_msg}"

        @tool
        def convert_result_to_natural_language_tool(sql_result: str, original_question: str) -> str:
            """将SQL查询结果转换为自然语言描述 - 增强版"""
            try:
                # 处理SQL执行失败的情况
                if "SQL执行失败" in sql_result:
                    return f"抱歉，查询「{original_question}」时遇到了问题。可能是数据库中没有相关数据，或者查询条件需要调整。请尝试重新描述您的需求。"

                # 提取数据信息
                import re
                import json

                # 改进的解析逻辑 - 支持多种格式
                data_match = re.search(r'数据：(\[.*\])', sql_result, re.DOTALL)
                row_count_match = re.search(r'查询结果：(\d+)行数据', sql_result)
                cols_match = re.search(r'列名：(\[.*?\])', sql_result)

                # 如果无法解析标准格式，尝试智能解析
                if not data_match or not row_count_match:
                    # 使用AI直接解析和转换
                    prompt = f"""请将以下数据库查询结果转换为自然、友好的中文描述：

用户问题：{original_question}
查询结果：{sql_result}

要求：
1. 用自然的中文描述查询结果
2. 如果是数字统计，要明确说明数量
3. 如果是列表数据，要概括主要信息
4. 保持简洁明了，避免技术术语
5. 语气要友好专业

请直接返回自然语言描述，不要包含任何SQL代码或技术细节。"""

                    response = self.model.invoke(prompt)
                    return response.content.strip()

                # 标准格式解析
                try:
                    # 使用更安全的JSON解析
                    data_str = data_match.group(1)
                    # 替换单引号为双引号以符合JSON格式
                    data_str = data_str.replace("'", '"').replace('None', 'null')
                    data = json.loads(data_str)
                    row_count = int(row_count_match.group(1))
                    cols = json.loads(cols_match.group(1).replace("'", '"')) if cols_match else []
                except Exception as parse_error:
                    logging.error(f"解析查询结果失败：{parse_error}")
                    # 降级到AI解析
                    prompt = f"""请将以下查询结果转换为自然语言：

用户问题：{original_question}
查询结果：{sql_result}

请用友好的中文描述结果，避免技术术语。"""

                    response = self.model.invoke(prompt)
                    return response.content.strip()

                # 生成自然语言描述
                if row_count == 0:
                    return f"根据您的查询「{original_question}」，数据库中暂时没有找到匹配的数据。您可以尝试调整查询条件或换个方式提问。"

                # 构建自然语言描述的提示
                data_preview = data[:3] if len(data) > 3 else data  # 减少预览数据量

                prompt = f"""请将以下SQL查询结果转换为自然、易懂的中文描述：

用户问题：{original_question}
查询结果数量：{row_count}行
主要字段：{cols}
数据示例：{data_preview}

要求：
1. 用自然的中文描述查询结果，语气友好专业
2. 如果数据较多，提供概要信息和关键统计
3. 突出重要信息，如数量、名称、特征等
4. 保持简洁明了，避免使用技术术语
5. 如果有多条记录，可以提及"等"来表示还有更多

请直接返回自然语言描述，不要包含SQL代码或技术细节。"""

                response = self.model.invoke(prompt)
                natural_description = response.content.strip()

                # 确保描述以友好的方式开始
                if not natural_description.startswith(("根据", "查询", "您", "数据库")):
                    natural_description = f"根据您的查询，{natural_description}"

                return natural_description

            except Exception as e:
                logging.error(f"结果转换失败：{str(e)}")
                return f"抱歉，在处理查询「{original_question}」的结果时遇到了技术问题。请稍后重试或换个方式提问。"

        @tool
        def fix_sql_tool(sql: str, error_message: str, context: str) -> str:
            """修复SQL语句中的错误"""
            prompt = f"""SQL语句执行失败，请修复错误：

原始SQL：{sql}
错误信息：{error_message}
数据库上下文：{context}

请分析错误原因并生成修正后的SQL语句。只返回修正后的SQL，不要包含解释。"""

            response = self.model.invoke(prompt)
            fixed_sql = self.extract_sql(response.content)
            return f"修复后的SQL语句：{fixed_sql}"

        @tool
        def diagnose_embedding_tool() -> str:
            """诊断嵌入API连接问题"""
            try:
                # 检查环境变量
                api_key = os.getenv("OPENAI_API_KEY")
                base_url = os.getenv("OPENAI_BASE_URL")
                model = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-large")

                diagnosis = f"""
环境变量检查：
- API_KEY: {'已设置' if api_key else '未设置'}
- BASE_URL: {base_url or '未设置'}
- MODEL: {model}

"""

                if not api_key:
                    return diagnosis + "错误：API密钥未设置"

                if not base_url:
                    return diagnosis + "错误：API地址未设置"

                # 测试简单的嵌入调用
                try:
                    test_embedding = self.embedding_fn.embed_query("test")
                    diagnosis += f"嵌入测试成功，向量维度：{len(test_embedding)}"
                    return diagnosis
                except Exception as e:
                    diagnosis += f"嵌入测试失败：{str(e)}"
                    return diagnosis

            except Exception as e:
                return f"诊断失败：{str(e)}"

        @tool
        def smart_query_tool(question: str) -> str:
            """智能查询工具 - 一站式处理用户问题，提升性能"""
            try:
                start_time = time.time()
                logging.info(f"开始智能查询: {question}")

                # 1. 并发检索相关信息
                retrieval_results = self.concurrent_retrieve_all(question)

                if "error" in retrieval_results:
                    return f"抱歉，查询「{question}」时遇到了数据检索问题。请稍后重试。"

                formatted = self.format_concurrent_results(retrieval_results)
                context = f"{formatted['ddl']}\n\n{formatted['examples']}\n\n{formatted['descriptions']}"

                # 2. 生成并执行SQL
                sql_prompt = f"""基于以下信息生成SQL语句：

{context}

用户问题：{question}

请生成准确的SQL语句。"""

                sql_response = self.model.invoke(sql_prompt)
                sql = self.extract_sql(sql_response.content)

                if not sql:
                    return f"抱歉，无法为问题「{question}」生成合适的查询。请尝试重新描述您的需求。"

                # 3. 执行SQL并处理结果
                try:
                    clean_sql = sql.strip().rstrip(';')
                    # 使用SQLAlchemy引擎执行SQL
                    with self.engine.connect() as connection:
                        result = connection.execute(text(clean_sql))
                        results = result.fetchall()
                        columns = list(result.keys())

                    if not results:
                        return f"查询「{question}」没有找到相关数据。数据库中可能暂时没有符合条件的记录。"

                    # 4. 智能转换为自然语言
                    result_data = [dict(zip(columns, row)) for row in results]

                    nl_prompt = f"""将以下查询结果转换为自然语言回答：

原始问题：{question}
查询结果：{result_data}

要求：
1. 用自然、友好的语言回答
2. 突出关键数据和数字
3. 如果结果很多，提供摘要
4. 不要显示SQL或技术细节
5. 直接回答用户问题"""

                    nl_response = self.model.invoke(nl_prompt)
                    final_answer = nl_response.content

                    # 记录性能
                    total_time = time.time() - start_time
                    logging.info(f"智能查询完成，总耗时：{total_time:.2f}秒")

                    return final_answer

                except Exception as sql_error:
                    # SQL执行失败，尝试修复一次
                    logging.warning(f"SQL执行失败，尝试修复: {sql_error}")

                    fix_prompt = f"""SQL执行失败，请修复：

原始SQL：{sql}
错误信息：{str(sql_error)}

请生成修复后的SQL。"""

                    try:
                        fix_response = self.model.invoke(fix_prompt)
                        fixed_sql = self.extract_sql(fix_response.content)

                        if fixed_sql and fixed_sql != sql:
                            with self.engine.connect() as connection:
                                result = connection.execute(text(fixed_sql.strip().rstrip(';')))
                                results = result.fetchall()
                                columns = list(result.keys())

                            if results:
                                result_data = [dict(zip(columns, row)) for row in results]
                                nl_prompt = f"""将查询结果转换为自然语言：

问题：{question}
结果：{result_data}

用友好的语言回答，不显示技术细节。"""

                                nl_response = self.model.invoke(nl_prompt)
                                return nl_response.content
                    except:
                        pass

                    return f"抱歉，查询「{question}」时遇到了技术问题。请尝试重新描述您的需求。"

            except Exception as e:
                logging.error(f"智能查询失败：{str(e)}")
                return f"抱歉，处理您的问题时出现了错误。请稍后重试或换个方式提问。"

        return [
            smart_query_tool,  # 新的一站式工具，优先使用
            diagnose_embedding_tool,
            concurrent_retrieve_all_tool,  # 🚀 优化版并发检索工具（集成统一嵌入）
            # 移除了 get_embedding_performance_tool - 简化版本不再提供性能统计
            generate_sql_tool,
            execute_sql_tool,
            convert_result_to_natural_language_tool,
            fix_sql_tool
        ]

    def extract_sql(self, text: str) -> str:
        """从文本中提取SQL语句"""
        # 尝试匹配 SQL 代码块
        sql_blocks = re.findall(r'```sql\n(.*?)\n```', text, re.DOTALL)
        if sql_blocks:
            return sql_blocks[0].strip()

        # 如果没有找到代码块，尝试匹配 SELECT 语句
        select_match = re.search(r'SELECT.*?;', text, re.DOTALL)
        if select_match:
            return select_match.group(0).strip()

        # 如果都没有找到，返回原始文本
        return text.strip()
    
    def create_workflow(self) -> StateGraph:
        """创建工作流程"""
        workflow = StateGraph(Dict[str, Any])
        
        # 添加节点
        workflow.add_node("agent", self.agent)
        
        # 设置入口和出口
        workflow.set_entry_point("agent")
        workflow.add_edge("agent", END)
        
        return workflow.compile()
    
    def process_query(self, question: str) -> Dict[str, Any]:
        """处理用户查询"""
        try:
            # 初始化状态
            state = {
                "messages": [HumanMessage(content=question)],
                "question": question
            }

            # 执行工作流程
            result = self.workflow.invoke(state)

            # 后处理：确保输出是自然语言而非SQL代码
            result = self._post_process_result(result, question)

            return result
        except Exception as e:
            logging.error(f"处理查询时出错：{e}")
            return {"error": str(e)}

    def _post_process_result(self, result: Dict[str, Any], original_question: str) -> Dict[str, Any]:
        """后处理结果，确保输出格式正确"""
        try:
            # 检查结果中是否包含消息
            if "messages" in result and result["messages"]:
                last_message = result["messages"][-1]
                content = last_message.content if hasattr(last_message, 'content') else str(last_message)

                # 检查是否包含SQL代码
                if self._contains_sql_code(content):
                    logging.warning("检测到输出包含SQL代码，进行自然语言转换")
                    # 使用AI将SQL结果转换为自然语言
                    natural_response = self._convert_sql_to_natural_language(content, original_question)

                    # 更新结果
                    from langchain_core.messages import AIMessage
                    result["messages"][-1] = AIMessage(content=natural_response)

            return result

        except Exception as e:
            logging.error(f"后处理结果时出错：{e}")
            return result

    def _contains_sql_code(self, content: str) -> bool:
        """检查内容是否包含SQL代码"""
        sql_indicators = [
            'SELECT', 'FROM', 'WHERE', 'JOIN', 'INSERT', 'UPDATE', 'DELETE',
            'CREATE', 'ALTER', 'DROP', 'GROUP BY', 'ORDER BY', 'HAVING'
        ]

        content_upper = content.upper()

        # 检查是否包含多个SQL关键词
        sql_count = sum(1 for indicator in sql_indicators if indicator in content_upper)

        # 如果包含2个或以上SQL关键词，且不是在自然语言描述中，则认为是SQL代码
        if sql_count >= 2:
            # 进一步检查是否是在代码块中
            if '<pre>' in content or '```' in content or 'SQL语句' in content:
                return True

        return False

    def _convert_sql_to_natural_language(self, content: str, original_question: str) -> str:
        """将包含SQL的内容转换为自然语言"""
        try:
            prompt = f"""以下内容包含了SQL代码，请将其转换为用户友好的自然语言描述：

用户问题：{original_question}
当前回复：{content}

要求：
1. 移除所有SQL代码和技术细节
2. 用自然的中文描述查询结果或操作
3. 保持友好专业的语气
4. 如果有具体数据，要清楚地说明
5. 避免使用技术术语

请直接返回自然语言描述："""

            response = self.model.invoke(prompt)
            natural_response = response.content.strip()

            # 确保响应不包含SQL代码
            if self._contains_sql_code(natural_response):
                return f"根据您的查询「{original_question}」，我已经在数据库中进行了搜索。如果您需要具体的结果信息，请告诉我您最关心的是哪些方面的数据。"

            return natural_response

        except Exception as e:
            logging.error(f"SQL转自然语言失败：{e}")
            return f"根据您的查询「{original_question}」，我已经处理了您的请求。如果您需要更详细的信息，请告诉我您想了解的具体内容。"

    def process_query_with_retry(self, question: str, max_retries: int = None) -> Dict[str, Any]:
        """处理用户查询，支持SQL错误重试修复"""
        if max_retries is None:
            max_retries = AGENT_MAX_RETRIES
        try:
            # 初始化状态
            state = {
                "messages": [HumanMessage(content=question)],
                "question": question,
                "retry_count": 0,
                "max_retries": max_retries
            }
            
            # 执行工作流程
            result = self.workflow.invoke(state)
            
            # 检查是否需要重试
            if "error" in result and "SQL执行失败" in str(result.get("error", "")):
                return self._retry_with_fix(question, result["error"], max_retries)
            
            return result
        except Exception as e:
            logging.error(f"处理查询时出错：{e}")
            return {"error": str(e)}
    
    def _retry_with_fix(self, question: str, error_msg: str, max_retries: int) -> Dict[str, Any]:
        """重试机制：修复SQL错误"""
        for attempt in range(max_retries):
            logging.info(f"[重试] 第 {attempt + 1} 次尝试修复SQL")
            
            # 重新检索上下文
            q_emb = self.embedding_fn.embed_query(question)
            
            # 检索DDL
            ddl_hits = self._retrieve_ddl(q_emb)
            ddl_context = "\n".join(hit.get("ddl_text", "") for hit in ddl_hits)
            
            # 检索示例
            example_hits = self._retrieve_examples(q_emb)
            example_context = "\n".join(
                f"问题：{hit.get('question', '')}\nSQL：{hit.get('sql_text', '')}\n"
                for hit in example_hits
            )
            
            # 检索字段描述
            desc_hits = self._retrieve_descriptions(q_emb)
            desc_context = "\n".join(
                f"{hit.get('table_name', '')}.{hit.get('column_name', '')}: {hit.get('description', '')}"
                for hit in desc_hits
            )
            
            # 组装完整上下文
            full_context = f"### Schema Definitions:\n{ddl_context}\n### Field Descriptions:\n{desc_context}\n### Examples:\n{example_context}"
            
            # 生成修复后的SQL
            fixed_sql = self._generate_fixed_sql(question, error_msg, full_context)
            
            # 执行修复后的SQL
            success, result = self._execute_sql(fixed_sql)
            
            if success:
                return {
                    "success": True,
                    "sql": fixed_sql,
                    "result": result,
                    "retry_count": attempt + 1
                }
            
            error_msg = result  # 更新错误信息用于下次重试
        
        return {
            "success": False,
            "error": f"经过 {max_retries} 次重试后仍然失败",
            "last_error": error_msg
        }
    
    def _retrieve_ddl(self, q_emb):
        """检索DDL信息"""
        try:
            self.milvus_client.load_collection("ddl_knowledge")
            results = self.milvus_client.search(
                collection_name="ddl_knowledge",
                data=[q_emb],
                limit=3,
                output_fields=["ddl_text"]
            )
            return results[0]
        except Exception as e:
            logging.error(f"检索DDL时出错：{e}")
            return []
    
    def _retrieve_examples(self, q_emb):
        """检索示例"""
        try:
            self.milvus_client.load_collection("q2sql_knowledge")
            results = self.milvus_client.search(
                collection_name="q2sql_knowledge",
                data=[q_emb],
                limit=3,
                output_fields=["question", "sql_text"]
            )
            return results[0]
        except Exception as e:
            logging.error(f"检索示例时出错：{e}")
            return []
    
    def _retrieve_descriptions(self, q_emb):
        """检索字段描述"""
        try:
            self.milvus_client.load_collection("dbdesc_knowledge")
            results = self.milvus_client.search(
                collection_name="dbdesc_knowledge",
                data=[q_emb],
                limit=5,
                output_fields=["table_name", "column_name", "description"]
            )
            return results[0]
        except Exception as e:
            logging.error(f"检索字段描述时出错：{e}")
            return []
    
    def _generate_fixed_sql(self, question: str, error_msg: str, context: str) -> str:
        """生成修复后的SQL"""
        prompt = f"""SQL执行失败，请修复错误：

数据库上下文：
{context}

用户问题：{question}
错误信息：{error_msg}

请分析错误原因并生成修正后的SQL语句。只返回SQL语句，不要包含任何解释。"""
        
        response = self.model.invoke(prompt)
        return self.extract_sql(response.content)
    
    def _execute_sql(self, sql: str) -> tuple[bool, str]:
        """执行SQL并返回结果（与execute_sql_tool保持一致）"""
        try:
            # 清理SQL语句
            clean_sql = sql.strip().rstrip(';')
            if not clean_sql:
                return False, "SQL语句为空"

            with self.engine.connect() as conn:
                result = conn.execute(text(clean_sql))
                cols = list(result.keys())
                rows = result.fetchall()

                # 转换为字典列表格式
                data = []
                for row in rows:
                    row_dict = {}
                    for i, col in enumerate(cols):
                        value = row[i]
                        if value is None:
                            row_dict[col] = None
                        elif isinstance(value, (int, float, str, bool)):
                            row_dict[col] = value
                        else:
                            row_dict[col] = str(value)
                    data.append(row_dict)

                result_str = f"查询结果：{len(data)}行数据\n列名：{cols}\n数据：{data}"
                return True, result_str

        except Exception as e:
            return False, str(e)

# 全局实例，方便其他类访问
text2sql_agent_instance = None

def get_text2sql_agent():
    """获取Text2SQL智能体的全局实例"""
    global text2sql_agent_instance
    if text2sql_agent_instance is None:
        text2sql_agent_instance = Text2SQLAgent()
    return text2sql_agent_instance

def main():
    """主函数"""
    agent = get_text2sql_agent()
    
    print("Text2SQL 智能体已启动！输入 'q' 退出")
    
    while True:
        try:
            question = input("\n请输入您的自然语言查询：")
            if question.lower() == 'q':
                break
            
            if question.strip():
                result = agent.process_query(question)
                print(f"\n处理结果：{result}")
        
        except KeyboardInterrupt:
            print("\n程序已退出")
            break
        except Exception as e:
            print(f"发生错误：{e}")

if __name__ == "__main__":
    main() 